#!/usr/bin/env python


from __future__ import with_statement

import errno
import os
import os.path
import re
import shlex
import signal
import sys

from optparse import OptionParser

DEBUG = False
TESTING = False
BASEPATH = None

class SievelogSyntaxError(Exception):
	"""Exception to be thrown when the sievelog syntax is wrong."""
	pass

class BasepathError(EnvironmentError):
	"""Exception type to be thrown when an output file is specified that's
	outside the BASEPATH."""
	pass

def memoize(func):
	memo_store = {}

	def wrap(*args):
		args = tuple(args)

		if args not in memo_store:
			memo_store[args] = func(*args)

		if DEBUG:
			print "memoize: dict:: %s" % memo_store
		return memo_store[args]

	return wrap

def warn(lineno, message):
	print >> sys.stderr, "warning: line %d: %s" % (lineno, message)

def regexify(guard):
	"""Create a compiled regex from the given guard statement."""

	flags = 0

	if guard.endswith('i'):
		flags |= re.IGNORECASE
		guard = guard[:-1]

	if (guard.startswith('"') and guard.endswith('"')):
		return re.compile(guard[1:-1], flags)
	elif guard == '*':
		return True
	else:
		raise SievelogSyntaxError(
			"guard '%s' must be enclosed in double-quotes. skipping..." % guard
		)

@memoize
def basic_actionize(action):
	"""Create a closure that takes the provided basic action."""

	if action == '/dev/null':
		# We can just directly drop on the floor messages bound for /dev/null
		return lambda x: None

	elif action.startswith('/'):
		# Wire up a closure that writes to the file named by `action'.
		try:
			if os.path.commonprefix((BASEPATH, action)) == BASEPATH:
				outf = file(action, 'a')
			else:
				raise BasepathError(
					"output file %s must be located under %s." %\
						(action, BASEPATH)
				)

		except (IOError, OSError, BasepathError), e:
			if TESTING:
				print "Error: %s" % e
				return lambda x: None
			else:
				raise
		def _write_outf(line):
			outf.write(line)
			outf.flush()
		return _write_outf

	elif action.startswith('mail'):
		# Wire up a closure that sends an email off.
		import smtplib
		from email.mime.text import MIMEText

		try:
			(subj, recipient, sender) = shlex.split(action)[1:]
		except ValueError:
			raise SievelogSyntaxError("bad mail action: %s" % action)

		def _mail_message(line):
			message = MIMEText(line)
			message['subject'] = subj
			message['to'] = recipient
			message['from'] = sender

			mailserver = smtplib.SMTP('localhost')
			mailserver.sendmail(sender, [recipient], message.as_string())

		return _mail_message
	else:
		raise SievelogSyntaxError(
			"unknown action type for '%s'. ignoring..." % action
		)

def parse_rule(sieverule):
	try:
		(guard, actions) = map(str.strip, sieverule.split("->", 1))
	except ValueError:
		raise SievelogSyntaxError(
			"syntax error in '%s'. skipping..." % sieverule
		)

	if actions.startswith('{'):
		# If the action begins with a brace it indicates we're
		# building a compound rule.
		#
		# HACK: In some sense this should be unified with basic_actionize(),
		# but I can't see a way to do that without that winding up being as
		# much of a hack as the existence of compound rules themselves are.

		# We're dealing with a bunch of '"regex" -> action' lines within
		# this compound rule.
		rules = actions.split('\n')[1:-1]

		# Transform the lines into a list of executable rules.
		rules = tuple(map(lambda sr: parse_rule(sr), rules))

		# Build the closure that will run through the rules in this
		# compound rule.
		def _do_rules(line):
			for (guard, action) in rules:	
				if guard.search(line):
					for fxn in action:
						fxn(line)
					break

		# Return the compound rule as just an ordinary executable rule.
		return (regexify(guard), (_do_rules,))
	else:
		# We have a basic rule, so we build its list of actions and
		# return the executable rule.
		actions = map(str.strip, actions.split(","))
		return (regexify(guard), tuple(map(basic_actionize, actions)))

def parse_file(sievefile):
	"""Lex and parse the rules file.

	The end result of this is `sieves' is a list populated with
	(guard, actions) pairs. `guard' is a boolean function of a message, and
	`actions' is a list of functions to be called on a message if it's matched
	by the guard.
	"""
	sieves = []
	line_acc = []
	mode = ""

	with open(sievefile) as sievedescs:
		for lineno, sievedesc in enumerate(sievedescs):
			sievedesc = sievedesc.strip()

			if not sievedesc or sievedesc.startswith('#'):
				continue

			if sievedesc.endswith('{'):
				mode = "IN_BRACES"

			if mode == "IN_BRACES":
				if sievedesc.endswith('}'):
					mode = ""
				else:
					if sievedesc.endswith('\\'):
						line_acc.append(sievedesc[:-1])
					else:
						line_acc.append(sievedesc + '\n')
					continue

			if sievedesc.endswith('\\'):
				line_acc.append(sievedesc[:-1])
				continue
			else:
				sievedesc = ' '.join(line_acc) + sievedesc
				sievedesc = sievedesc.strip()
				line_acc = []

			try:
				sieves.append(parse_rule(sievedesc))
			except SievelogSyntaxError, e:
				warn(lineno + 1, e)

		if mode == "IN_BRACES":
			print >> sys.stderr, "parse error: braced statement never closed."
			sys.exit(1)

	return sieves

def read_stdin():
	"""Provide an iterator that yields lines of stdin and ignores EINTRs caused
	by the user sending SIGINFOs."""
	while True:
		try:
			yield sys.stdin.next()
		except IOError, ex:
			if ex.errno == errno.EINTR:
				continue
			else:
				raise

def engine(sieves, counters):
	"""Run the sieve rules."""

	# We handle catchalls differently from ordinary sieve rules.
	catchalls = set(reduce(list.extend, [a for (g, a) in sieves if g == True]))
	sieves = [(g, a) for (g, a) in sieves if g != True]

	for rule in sieves:
		counters[rule] = 0

	for message in read_stdin():
		for (index, (guard, actions)) in enumerate(sieves):
			if guard.search(message):
				for action in actions:
					action(message)

				if opts.adaptive:
					# Update our counter table.
					# If this rule has been matched more than the one before
					# it in the list, we swap them.
					rule = (guard, actions)
					counters[rule] += 1
					if index > 0 and counters[rule] > counters[sieves[index-1]]:
						(sieves[index], sieves[index-1]) = \
							(sieves[index-1], sieves[index])
				break
		else:
			for action in catchalls:
				action(message)

def output_stats(ctr_table, fobj):
	"""Output the counter table to the provided file-like object."""

	fobj.seek(0)
	fobj.truncate()

	stats = [(grd, cnt) for ((grd, act), cnt) in ctr_table.iteritems()]
	for (guard, count) in sorted(stats, key = lambda x: x[1]):
		fobj.write("%d '%s'\n" % (count, guard.pattern))
	fobj.flush()

def main(opts, sievefile):
	sieves = parse_file(sievefile)

	counters = {}

	if opts.stats:
		# Wire up the SIGINFO handler to dump out the counter table
		stats_file = file(opts.stats, 'w')

		def handle_siginfo(sig, frame):
			output_stats(counters, stats_file)

		signal.signal(signal.SIGINFO, handle_siginfo)

	if not TESTING:
		engine(sieves, counters)

	
if __name__ == '__main__':

	def set_debug(*_):
		global DEBUG
		DEBUG = True

	def set_testing(*_):
		global TESTING
		TESTING = True

	op = OptionParser("usage: %prog [options] <sievefile>")

	op.add_option("-t", "--test", action="callback", callback=set_testing,
		help="Only do a syntax check of the sievefile.")
	op.add_option("-D", "--debug", action="callback", callback=set_debug,
		help="Turn on debug mode.")
	op.add_option("-o", "--optimize", dest="adaptive", default=False,
		action="store_true",
		help="Turn on the adaptive optimizer for the sieve ruleset")
	op.add_option("-P", "--stats", dest="stats", default="", metavar="FILE",
		help="Write statistics to FILE on receipt of SIGINFO")
	op.add_option("-R", "--restrict", dest="basepath", default="/",
		help="Restrict output files to within BASEPATH.")

	(opts, args) = op.parse_args()

	try:
		(sievefile,) = args
	except ValueError:
		op.print_help(sys.stderr)
		sys.exit(os.EX_USAGE)

	if opts.basepath.endswith(os.path.sep):
		BASEPATH = opts.basepath
	else:
		BASEPATH = opts.basepath + os.path.sep
		
	main(opts, sievefile)
